package com.study.chapter09;

import com.study.chapter05.source.ClickSource;
import com.study.entity.Event;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.*;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * @Description:
 * @Author: LiuQun
 * @Date: 2022/8/15 21:38
 */
public class KeyedStateTest {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                                .withTimestampAssigner((event, l) -> event.timestamp)
                );

        stream.keyBy(event -> event.user)
                .flatMap(new MyFlatMapFunc())
                .print();

        env.execute();
    }

    //自定义flatMapFunction
    public static class MyFlatMapFunc extends RichFlatMapFunction<Event,String>{
        //ValueState
        private ValueState<Event> valueState;
        private ListState<Event> listState;
        private MapState<String,Long> mapState;
        private ReducingState<Event> reducingState;
        private AggregatingState<String,String> aggregatingState;


        @Override
        public void open(Configuration parameters) throws Exception {
            //定义状态
            valueState = getRuntimeContext().getState(
                    new ValueStateDescriptor<Event>("my-value-event",Event.class)
            );

            listState = getRuntimeContext().getListState(
                    new ListStateDescriptor<Event>("my-list-event",Event.class)
            );
            mapState = getRuntimeContext().getMapState(
                    new MapStateDescriptor<String, Long>("my-map-event",String.class,Long.class)
            );
            reducingState = getRuntimeContext().getReducingState(
                    new ReducingStateDescriptor<Event>("my-reduce-event", new ReduceFunction<Event>() {
                        @Override
                        public Event reduce(Event value1, Event value2) throws Exception {
                            return new Event(value1.user,value1.url,value2.timestamp);
                        }
                    }, Event.class)
            );
            aggregatingState = getRuntimeContext().getAggregatingState(
                    new AggregatingStateDescriptor<String, Long, String>("my-agg-event", new AggregateFunction<String, Long, String>() {

                        @Override
                        public Long createAccumulator() {
                            return 0L;
                        }

                        @Override
                        public Long add(String value, Long accumulator) {
                            return accumulator + 1;
                        }

                        @Override
                        public String getResult(Long accumulator) {
                            return "count："+accumulator;
                        }

                        @Override
                        public Long merge(Long a, Long b) {
                            return a+b;
                        }
                    }, Long.class)
            );

        }

        @Override
        public void flatMap(Event value, Collector<String> out) throws Exception {
            // System.out.println(valueState.value());
            // valueState.update(value);
            // System.out.println("My Value："+valueState.value());

            listState.add(value);

            mapState.put(value.user,mapState.get(value.user) == null ? 1 : mapState.get(value.user) + 1);
            System.out.println("My Map State："+value.user +" "+mapState.get(value.user));

            reducingState.add(value);
            System.out.println("My Reduce State："+reducingState.get());

            aggregatingState.add(value.user);
            System.out.println("My Agg State："+aggregatingState.get());
        }
    }
}
